{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71b51443",
   "metadata": {},
   "source": [
    "# Fed-MVKM Advanced Visualization and Diagnostic Tools\n",
    "\n",
    "This notebook provides advanced visualization and diagnostic tools for analyzing the Fed-MVKM algorithm's behavior on the DHA dataset.\n",
    "\n",
    "**Author:** Kristina P. Sinaga  \n",
    "**Date:** May 2024  \n",
    "**Version:** 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a7836e",
   "metadata": {},
   "source": [
    "## 1. Setup and Imports\n",
    "\n",
    "First, let's import necessary libraries and set up our visualization environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c57447e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from pathlib import Path\n",
    "import scipy.io as sio\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.decomposition import PCA\n",
    "import plotly.express as px\n",
    "import plotly.graph_objects as go\n",
    "from plotly.subplots import make_subplots\n",
    "import pandas as pd\n",
    "from typing import List, Dict, Tuple\n",
    "import networkx as nx\n",
    "from scipy.stats import entropy\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "\n",
    "# Import our Fed-MVKM implementation\n",
    "from mvkm_ed import FedMVKMED, FedMVKMEDConfig\n",
    "from mvkm_ed.utils import MVKMEDDataProcessor, MVKMEDMetrics, MVKMEDVisualizer\n",
    "\n",
    "# Set style for static plots\n",
    "plt.style.use('seaborn')\n",
    "sns.set_palette(\"husl\")\n",
    "\n",
    "# Configure notebook for plotly\n",
    "import plotly.io as pio\n",
    "pio.templates.default = \"plotly_white\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1080fb4b",
   "metadata": {},
   "source": [
    "## 2. Advanced Visualization Classes\n",
    "\n",
    "Let's create specialized classes for different types of visualizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8d2f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClusterVisualization:\n",
    "    \"\"\"Advanced cluster visualization tools.\"\"\"\n",
    "    \n",
    "    def __init__(self, views: List[np.ndarray], labels: np.ndarray, \n",
    "                 predictions: np.ndarray):\n",
    "        self.views = views\n",
    "        self.labels = labels\n",
    "        self.predictions = predictions\n",
    "        self.n_views = len(views)\n",
    "        \n",
    "    def plot_cluster_embeddings(self, method='tsne', interactive=True):\n",
    "        \"\"\"Plot cluster embeddings using dimensionality reduction.\"\"\"\n",
    "        fig = make_subplots(rows=1, cols=self.n_views, \n",
    "                           subplot_titles=[f'View {i+1}' for i in range(self.n_views)])\n",
    "        \n",
    "        for i, view in enumerate(self.views):\n",
    "            # Dimensionality reduction\n",
    "            if method.lower() == 'tsne':\n",
    "                embedding = TSNE(n_components=2, random_state=42).fit_transform(view)\n",
    "            else:\n",
    "                embedding = PCA(n_components=2, random_state=42).fit_transform(view)\n",
    "            \n",
    "            # Create DataFrame for plotting\n",
    "            df = pd.DataFrame({\n",
    "                'x': embedding[:, 0],\n",
    "                'y': embedding[:, 1],\n",
    "                'True Label': self.labels,\n",
    "                'Predicted': self.predictions\n",
    "            })\n",
    "            \n",
    "            if interactive:\n",
    "                # Plotly interactive plot\n",
    "                scatter1 = go.Scatter(\n",
    "                    x=df['x'], y=df['y'],\n",
    "                    mode='markers',\n",
    "                    marker=dict(color=df['True Label'], \n",
    "                              showscale=True,\n",
    "                              colorscale='Viridis'),\n",
    "                    name='True Labels',\n",
    "                    showlegend=i==0\n",
    "                )\n",
    "                scatter2 = go.Scatter(\n",
    "                    x=df['x'], y=df['y'],\n",
    "                    mode='markers',\n",
    "                    marker=dict(color=df['Predicted'],\n",
    "                              showscale=True,\n",
    "                              colorscale='Plasma'),\n",
    "                    name='Predictions',\n",
    "                    visible='legendonly',\n",
    "                    showlegend=i==0\n",
    "                )\n",
    "                fig.add_trace(scatter1, row=1, col=i+1)\n",
    "                fig.add_trace(scatter2, row=1, col=i+1)\n",
    "            else:\n",
    "                # Matplotlib static plot\n",
    "                plt.figure(figsize=(6, 6))\n",
    "                plt.scatter(df['x'], df['y'], c=df['True Label'], \n",
    "                          cmap='viridis', alpha=0.6)\n",
    "                plt.title(f'View {i+1} Clusters')\n",
    "                plt.colorbar(label='Cluster')\n",
    "                plt.show()\n",
    "        \n",
    "        if interactive:\n",
    "            fig.update_layout(height=500, width=250*self.n_views,\n",
    "                            title=f'Cluster Embeddings ({method.upper()})')\n",
    "            fig.show()\n",
    "            \n",
    "    def plot_view_relationships(self):\n",
    "        \"\"\"Visualize relationships between views using correlation network.\"\"\"\n",
    "        n_views = len(self.views)\n",
    "        correlations = np.zeros((n_views, n_views))\n",
    "        \n",
    "        # Compute correlations between views\n",
    "        for i in range(n_views):\n",
    "            for j in range(n_views):\n",
    "                if i != j:\n",
    "                    corr = np.corrcoef(\n",
    "                        self.views[i].reshape(-1),\n",
    "                        self.views[j].reshape(-1)\n",
    "                    )[0, 1]\n",
    "                    correlations[i, j] = abs(corr)\n",
    "        \n",
    "        # Create network graph\n",
    "        G = nx.Graph()\n",
    "        for i in range(n_views):\n",
    "            G.add_node(f'View {i+1}')\n",
    "        \n",
    "        for i in range(n_views):\n",
    "            for j in range(i+1, n_views):\n",
    "                if correlations[i, j] > 0.1:  # Threshold for visibility\n",
    "                    G.add_edge(f'View {i+1}', f'View {j+1}', \n",
    "                             weight=correlations[i, j])\n",
    "        \n",
    "        plt.figure(figsize=(8, 8))\n",
    "        pos = nx.spring_layout(G)\n",
    "        \n",
    "        # Draw nodes\n",
    "        nx.draw_networkx_nodes(G, pos, node_color='lightblue', \n",
    "                             node_size=1000, alpha=0.7)\n",
    "        \n",
    "        # Draw edges with varying thickness\n",
    "        edges = G.edges()\n",
    "        weights = [G[u][v]['weight'] * 5 for u, v in edges]\n",
    "        nx.draw_networkx_edges(G, pos, width=weights, alpha=0.5)\n",
    "        \n",
    "        # Add labels\n",
    "        nx.draw_networkx_labels(G, pos)\n",
    "        \n",
    "        plt.title('View Relationship Network\\n(Edge thickness indicates correlation strength)')\n",
    "        plt.axis('off')\n",
    "        plt.tight_layout()\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76c425dd",
   "metadata": {},
   "source": [
    "## 3. Federation Analysis Tools\n",
    "\n",
    "Tools for analyzing the federated learning process and client behavior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec0606d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FederationAnalyzer:\n",
    "    \"\"\"Tools for analyzing federated learning behavior.\"\"\"\n",
    "    \n",
    "    def __init__(self, model: FedMVKMED, client_data: Dict, client_labels: Dict):\n",
    "        self.model = model\n",
    "        self.client_data = client_data\n",
    "        self.client_labels = client_labels\n",
    "        self.history = model.history\n",
    "        \n",
    "    def plot_client_convergence_3d(self):\n",
    "        \"\"\"Create 3D visualization of client convergence trajectories.\"\"\"\n",
    "        fig = go.Figure()\n",
    "        \n",
    "        for client_id, objectives in self.history['client_objectives'].items():\n",
    "            fig.add_trace(go.Scatter3d(\n",
    "                x=np.arange(len(objectives)),\n",
    "                y=[client_id] * len(objectives),\n",
    "                z=objectives,\n",
    "                mode='lines+markers',\n",
    "                name=f'Client {client_id}',\n",
    "                line=dict(width=4),\n",
    "                marker=dict(size=4)\n",
    "            ))\n",
    "        \n",
    "        fig.update_layout(\n",
    "            title='3D Client Convergence Trajectories',\n",
    "            scene=dict(\n",
    "                xaxis_title='Iteration',\n",
    "                yaxis_title='Client ID',\n",
    "                zaxis_title='Objective Value',\n",
    "                camera=dict(\n",
    "                    up=dict(x=0, y=0, z=1),\n",
    "                    center=dict(x=0, y=0, z=0),\n",
    "                    eye=dict(x=1.5, y=1.5, z=1.5)\n",
    "                )\n",
    "            ),\n",
    "            width=800,\n",
    "            height=800\n",
    "        )\n",
    "        \n",
    "        fig.show()\n",
    "        \n",
    "    def analyze_client_diversity(self):\n",
    "        \"\"\"Analyze and visualize diversity across clients.\"\"\"\n",
    "        n_clients = len(self.client_data)\n",
    "        \n",
    "        # Compute statistics for each client\n",
    "        stats = []\n",
    "        for client_id in self.client_data:\n",
    "            client_views = self.client_data[client_id]\n",
    "            client_labels = self.client_labels[client_id]\n",
    "            \n",
    "            # Basic statistics\n",
    "            n_samples = len(client_labels)\n",
    "            n_classes = len(np.unique(client_labels))\n",
    "            class_entropy = entropy(\n",
    "                np.unique(client_labels, return_counts=True)[1] / n_samples\n",
    "            )\n",
    "            \n",
    "            # View statistics\n",
    "            view_stats = []\n",
    "            for view in client_views:\n",
    "                mean_norm = np.linalg.norm(view.mean(axis=0))\n",
    "                std_norm = np.linalg.norm(view.std(axis=0))\n",
    "                view_stats.extend([mean_norm, std_norm])\n",
    "            \n",
    "            stats.append([client_id, n_samples, n_classes, class_entropy] + view_stats)\n",
    "        \n",
    "        # Create visualization\n",
    "        fig = make_subplots(\n",
    "            rows=2, cols=2,\n",
    "            subplot_titles=(\n",
    "                'Sample Distribution',\n",
    "                'Class Distribution',\n",
    "                'Feature Statistics',\n",
    "                'Class Entropy'\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        # Sample distribution\n",
    "        fig.add_trace(\n",
    "            go.Bar(x=[f'Client {s[0]}' for s in stats],\n",
    "                  y=[s[1] for s in stats],\n",
    "                  name='Samples'),\n",
    "            row=1, col=1\n",
    "        )\n",
    "        \n",
    "        # Class distribution\n",
    "        fig.add_trace(\n",
    "            go.Bar(x=[f'Client {s[0]}' for s in stats],\n",
    "                  y=[s[2] for s in stats],\n",
    "                  name='Classes'),\n",
    "            row=1, col=2\n",
    "        )\n",
    "        \n",
    "        # Feature statistics\n",
    "        fig.add_trace(\n",
    "            go.Box(y=[s[4:] for s in stats],\n",
    "                  name='View Stats',\n",
    "                  boxpoints='all'),\n",
    "            row=2, col=1\n",
    "        )\n",
    "        \n",
    "        # Class entropy\n",
    "        fig.add_trace(\n",
    "            go.Scatter(x=[f'Client {s[0]}' for s in stats],\n",
    "                      y=[s[3] for s in stats],\n",
    "                      mode='lines+markers',\n",
    "                      name='Entropy'),\n",
    "            row=2, col=2\n",
    "        )\n",
    "        \n",
    "        fig.update_layout(height=800, width=1000,\n",
    "                         showlegend=False,\n",
    "                         title_text=\"Client Diversity Analysis\")\n",
    "        fig.show()\n",
    "        \n",
    "    def plot_model_evolution(self):\n",
    "        \"\"\"Visualize the evolution of model parameters over iterations.\"\"\"\n",
    "        if 'center_updates' not in self.history:\n",
    "            print(\"No center updates found in history\")\n",
    "            return\n",
    "            \n",
    "        center_updates = self.history['center_updates']\n",
    "        n_iterations = len(center_updates)\n",
    "        n_views = len(center_updates[0])\n",
    "        \n",
    "        fig = make_subplots(\n",
    "            rows=n_views, cols=1,\n",
    "            subplot_titles=[f'View {i+1} Center Evolution' for i in range(n_views)],\n",
    "            vertical_spacing=0.1\n",
    "        )\n",
    "        \n",
    "        for view_idx in range(n_views):\n",
    "            # Get center trajectories for this view\n",
    "            centers = np.array([update[view_idx] for update in center_updates])\n",
    "            \n",
    "            # Create heatmap of center movement\n",
    "            fig.add_trace(\n",
    "                go.Heatmap(\n",
    "                    z=centers.reshape(n_iterations, -1),\n",
    "                    colorscale='Viridis',\n",
    "                    showscale=True,\n",
    "                    name=f'View {view_idx+1}'\n",
    "                ),\n",
    "                row=view_idx+1, col=1\n",
    "            )\n",
    "        \n",
    "        fig.update_layout(\n",
    "            height=300*n_views,\n",
    "            width=800,\n",
    "            title_text=\"Model Parameter Evolution\"\n",
    "        )\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e42f5d",
   "metadata": {},
   "source": [
    "## 4. Performance Diagnostic Tools\n",
    "\n",
    "Tools for detailed performance analysis and debugging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8bb8e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PerformanceDiagnostics:\n",
    "    \"\"\"Advanced diagnostic tools for Fed-MVKM performance analysis.\"\"\"\n",
    "    \n",
    "    def __init__(self, model: FedMVKMED, views: List[np.ndarray], \n",
    "                 predictions: np.ndarray, true_labels: np.ndarray):\n",
    "        self.model = model\n",
    "        self.views = views\n",
    "        self.predictions = predictions\n",
    "        self.true_labels = true_labels\n",
    "        \n",
    "    def analyze_cluster_quality(self):\n",
    "        \"\"\"Analyze cluster quality using multiple metrics.\"\"\"\n",
    "        from sklearn.metrics import silhouette_samples, calinski_harabasz_score\n",
    "        \n",
    "        results = {}\n",
    "        fig = make_subplots(\n",
    "            rows=2, cols=2,\n",
    "            subplot_titles=(\n",
    "                'Silhouette Scores Distribution',\n",
    "                'Cluster Sizes',\n",
    "                'Intra-cluster Distances',\n",
    "                'Inter-cluster Distances'\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        for view_idx, view in enumerate(self.views):\n",
    "            # Silhouette scores\n",
    "            sil_scores = silhouette_samples(view, self.predictions)\n",
    "            fig.add_trace(\n",
    "                go.Violin(\n",
    "                    y=sil_scores,\n",
    "                    name=f'View {view_idx+1}',\n",
    "                    box_visible=True,\n",
    "                    meanline_visible=True\n",
    "                ),\n",
    "                row=1, col=1\n",
    "            )\n",
    "            \n",
    "            # Cluster sizes\n",
    "            unique, counts = np.unique(self.predictions, return_counts=True)\n",
    "            fig.add_trace(\n",
    "                go.Bar(\n",
    "                    x=[f'Cluster {i}' for i in unique],\n",
    "                    y=counts,\n",
    "                    name=f'View {view_idx+1}'\n",
    "                ),\n",
    "                row=1, col=2\n",
    "            )\n",
    "            \n",
    "            # Intra-cluster distances\n",
    "            intra_dist = []\n",
    "            for cluster in unique:\n",
    "                mask = self.predictions == cluster\n",
    "                if np.sum(mask) > 1:\n",
    "                    cluster_points = view[mask]\n",
    "                    distances = pdist(cluster_points)\n",
    "                    intra_dist.extend(distances)\n",
    "            \n",
    "            fig.add_trace(\n",
    "                go.Box(\n",
    "                    y=intra_dist,\n",
    "                    name=f'View {view_idx+1}',\n",
    "                    boxpoints='outliers'\n",
    "                ),\n",
    "                row=2, col=1\n",
    "            )\n",
    "            \n",
    "            # Inter-cluster distances\n",
    "            centers = []\n",
    "            for cluster in unique:\n",
    "                mask = self.predictions == cluster\n",
    "                centers.append(view[mask].mean(axis=0))\n",
    "            center_dist = pdist(centers)\n",
    "            \n",
    "            fig.add_trace(\n",
    "                go.Box(\n",
    "                    y=center_dist,\n",
    "                    name=f'View {view_idx+1}',\n",
    "                    boxpoints='all'\n",
    "                ),\n",
    "                row=2, col=2\n",
    "            )\n",
    "            \n",
    "            # Store numerical results\n",
    "            results[f'view_{view_idx+1}'] = {\n",
    "                'silhouette_mean': np.mean(sil_scores),\n",
    "                'calinski_harabasz': calinski_harabasz_score(view, self.predictions),\n",
    "                'cluster_sizes': dict(zip(unique, counts)),\n",
    "                'intra_cluster_dist_mean': np.mean(intra_dist),\n",
    "                'inter_cluster_dist_mean': np.mean(center_dist)\n",
    "            }\n",
    "        \n",
    "        fig.update_layout(height=800, width=1000,\n",
    "                         title_text=\"Cluster Quality Analysis\")\n",
    "        fig.show()\n",
    "        \n",
    "        return results\n",
    "    \n",
    "    def analyze_view_importance(self):\n",
    "        \"\"\"Analyze the importance and contribution of each view.\"\"\"\n",
    "        if not hasattr(self.model, 'history') or 'view_weights' not in self.model.history:\n",
    "            print(\"No view weight history available\")\n",
    "            return\n",
    "            \n",
    "        view_weights = np.array(self.model.history['view_weights'])\n",
    "        n_iterations, n_views = view_weights.shape\n",
    "        \n",
    "        fig = make_subplots(\n",
    "            rows=2, cols=2,\n",
    "            subplot_titles=(\n",
    "                'View Weight Evolution',\n",
    "                'Final Weight Distribution',\n",
    "                'Weight Stability',\n",
    "                'View Correlation'\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        # View weight evolution\n",
    "        for i in range(n_views):\n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    y=view_weights[:, i],\n",
    "                    mode='lines',\n",
    "                    name=f'View {i+1}'\n",
    "                ),\n",
    "                row=1, col=1\n",
    "            )\n",
    "        \n",
    "        # Final weight distribution\n",
    "        fig.add_trace(\n",
    "            go.Pie(\n",
    "                labels=[f'View {i+1}' for i in range(n_views)],\n",
    "                values=view_weights[-1],\n",
    "                hole=0.3\n",
    "            ),\n",
    "            row=1, col=2\n",
    "        )\n",
    "        \n",
    "        # Weight stability\n",
    "        weight_std = view_weights.std(axis=0)\n",
    "        fig.add_trace(\n",
    "            go.Bar(\n",
    "                x=[f'View {i+1}' for i in range(n_views)],\n",
    "                y=weight_std,\n",
    "                name='Weight Stability'\n",
    "            ),\n",
    "            row=2, col=1\n",
    "        )\n",
    "        \n",
    "        # View correlation\n",
    "        corr_matrix = np.corrcoef([view.reshape(-1) for view in self.views])\n",
    "        fig.add_trace(\n",
    "            go.Heatmap(\n",
    "                z=corr_matrix,\n",
    "                x=[f'View {i+1}' for i in range(n_views)],\n",
    "                y=[f'View {i+1}' for i in range(n_views)],\n",
    "                colorscale='RdBu'\n",
    "            ),\n",
    "            row=2, col=2\n",
    "        )\n",
    "        \n",
    "        fig.update_layout(height=800, width=1000,\n",
    "                         title_text=\"View Importance Analysis\")\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dde61859",
   "metadata": {},
   "source": [
    "## 5. Example Usage\n",
    "\n",
    "Let's demonstrate how to use these visualization and diagnostic tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7595e77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data and model from previous analysis\n",
    "# Assuming you have run the main tutorial notebook first\n",
    "data_dir = Path(\"../data\")\n",
    "\n",
    "def load_and_preprocess():\n",
    "    \"\"\"Load and preprocess the DHA dataset.\"\"\"\n",
    "    processor = MVKMEDDataProcessor()\n",
    "    rgb_data = sio.loadmat(data_dir / 'RGB_DHA.mat')['RGB_DHA']\n",
    "    depth_data = sio.loadmat(data_dir / 'Depth_DHA.mat')['Depth_DHA']\n",
    "    labels = sio.loadmat(data_dir / 'label_DHA.mat')['label_DHA'].ravel()\n",
    "    \n",
    "    views = processor.preprocess_views([rgb_data, depth_data])\n",
    "    return views, labels\n",
    "\n",
    "# Load data\n",
    "views, true_labels = load_and_preprocess()\n",
    "\n",
    "# Create visualizers\n",
    "cluster_viz = ClusterVisualization(views, true_labels, model.predict(client_data))\n",
    "fed_analyzer = FederationAnalyzer(model, client_data, client_labels)\n",
    "diagnostics = PerformanceDiagnostics(model, views, all_predictions, all_true_labels)\n",
    "\n",
    "# Generate visualizations\n",
    "print(\"1. Cluster Embeddings:\")\n",
    "cluster_viz.plot_cluster_embeddings(method='tsne', interactive=True)\n",
    "\n",
    "print(\"\\n2. View Relationships:\")\n",
    "cluster_viz.plot_view_relationships()\n",
    "\n",
    "print(\"\\n3. Federation Analysis:\")\n",
    "fed_analyzer.plot_client_convergence_3d()\n",
    "fed_analyzer.analyze_client_diversity()\n",
    "fed_analyzer.plot_model_evolution()\n",
    "\n",
    "print(\"\\n4. Performance Diagnostics:\")\n",
    "cluster_quality = diagnostics.analyze_cluster_quality()\n",
    "diagnostics.analyze_view_importance()\n",
    "\n",
    "# Print numerical results\n",
    "print(\"\\nCluster Quality Metrics:\")\n",
    "print(json.dumps(cluster_quality, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4796284b",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook provided advanced visualization and diagnostic tools for:\n",
    "1. Cluster analysis and embedding visualization\n",
    "2. Federation behavior analysis\n",
    "3. Client diversity and contribution analysis\n",
    "4. Performance diagnostics and quality metrics\n",
    "\n",
    "These tools help in understanding and improving the Fed-MVKM algorithm's behavior on multi-view datasets."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
